"""Unit tests for the SkyPilot API server."""

import argparse
import asyncio
import threading
import time
from unittest import mock

import fastapi
import pytest
import uvicorn

from sky import models
from sky.server import server
from sky.server.requests import executor
from sky.skylet import constants
from sky.utils import common_utils
from sky.utils import config_utils


@mock.patch('uvicorn.run')
@mock.patch('sky.server.requests.executor.start')
@mock.patch('sky.utils.common_utils.get_cpu_count')
def test_deploy_flag_sets_workers_to_cpu_count(mock_get_cpu_count,
                                               mock_executor_start,
                                               mock_uvicorn_run):
    """Test that --deploy flag sets workers to CPU count."""
    # Setup
    mock_get_cpu_count.return_value = 8
    mock_executor_start.return_value = []

    # Create mock args with deploy=True
    test_args = argparse.Namespace(host='127.0.0.1', port=46580, deploy=True)

    # Call the main block with mocked args
    with mock.patch('argparse.ArgumentParser.parse_args',
                    return_value=test_args):
        with mock.patch('sky.server.requests.requests.reset_db_and_logs'):
            with mock.patch('sky.usage.usage_lib.maybe_show_privacy_policy'):
                # Execute the main block code directly
                num_workers = None
                if test_args.deploy:
                    num_workers = mock_get_cpu_count()

                workers = []
                try:
                    workers = mock_executor_start(test_args.deploy)
                    uvicorn.run('sky.server.server:app',
                                host=test_args.host,
                                port=test_args.port,
                                workers=num_workers)
                except Exception:
                    pass
                finally:
                    for worker in workers:
                        worker.terminate()

    # Verify that uvicorn.run was called with the correct number of workers
    mock_uvicorn_run.assert_called_once()
    call_args = mock_uvicorn_run.call_args[1]
    assert call_args['workers'] == 8
    assert call_args['host'] == '127.0.0.1'
    assert call_args['port'] == 46580


@mock.patch('uvicorn.run')
@mock.patch('sky.server.requests.executor.start')
def test_no_deploy_flag_uses_default_workers(mock_executor_start,
                                             mock_uvicorn_run):
    """Test that without --deploy flag, workers is None (default)."""
    # Setup
    mock_executor_start.return_value = []

    # Create mock args with deploy=False
    test_args = argparse.Namespace(host='127.0.0.1', port=46580, deploy=False)

    # Call the main block with mocked args
    with mock.patch('argparse.ArgumentParser.parse_args',
                    return_value=test_args):
        with mock.patch('sky.server.requests.requests.reset_db_and_logs'):
            with mock.patch('sky.usage.usage_lib.maybe_show_privacy_policy'):
                # Execute the main block code directly
                num_workers = None
                if test_args.deploy:
                    num_workers = common_utils.get_cpu_count()

                workers = []
                try:
                    workers = mock_executor_start(test_args.deploy)
                    uvicorn.run('sky.server.server:app',
                                host=test_args.host,
                                port=test_args.port,
                                workers=num_workers)
                except Exception:
                    pass
                finally:
                    for worker in workers:
                        worker.terminate()

    # Verify that uvicorn.run was called with workers=None
    mock_uvicorn_run.assert_called_once()
    call_args = mock_uvicorn_run.call_args[1]
    assert call_args['workers'] is None
    assert call_args['host'] == '127.0.0.1'
    assert call_args['port'] == 46580


@pytest.mark.asyncio
async def test_validate():
    """Test the validate endpoint."""
    mock_dag = mock.MagicMock()
    mock_validate_body = mock.MagicMock()
    mock_validate_body.dag = 'test_dag_yaml'
    mock_validate_body.request_options = {}

    with mock.patch('sky.server.server.dag_utils.load_chain_dag_from_yaml_str',
                   return_value=mock_dag), \
         mock.patch('sky.server.server.admin_policy_utils.apply',
                   return_value=(mock_dag, config_utils.Config())), \
         mock.patch.object(mock_dag, 'validate') as mock_validate:
        # Call validate endpoint
        await server.validate(mock_validate_body)
        # Verify validate was called with correct args
        mock_validate.assert_called_once_with(skip_file_mounts=True,
                                              skip_workdir=True)

    error_msg = 'Invalid DAG'
    with mock.patch('sky.server.server.dag_utils.load_chain_dag_from_yaml_str',
                   return_value=mock_dag), \
         mock.patch('sky.server.server.admin_policy_utils.apply',
                   return_value=(mock_dag, config_utils.Config())), \
         mock.patch.object(mock_dag, 'validate',
                          side_effect=ValueError(error_msg)):
        with pytest.raises(fastapi.HTTPException) as exc_info:
            await server.validate(mock_validate_body)
        assert exc_info.value.status_code == 400
        assert error_msg in str(exc_info.value.detail)

    # Create an event to track when validation completes
    validation_complete = asyncio.Event()

    def slow_validate(*args, **kwargs):
        # Simulate slow validation
        time.sleep(0.1)
        validation_complete.set()

    with mock.patch('sky.server.server.dag_utils.load_chain_dag_from_yaml_str',
                   return_value=mock_dag), \
         mock.patch('sky.server.server.admin_policy_utils.apply',
                   return_value=(mock_dag, config_utils.Config())), \
         mock.patch.object(mock_dag, 'validate',
                          side_effect=slow_validate):
        # Start validation in background
        validation_task = asyncio.create_task(
            server.validate(mock_validate_body))

        # Check that validation hasn't completed immediately
        assert not validation_complete.is_set()

        # Wait for validation to complete
        await validation_task
        assert validation_complete.is_set()


@pytest.mark.asyncio
async def test_logs():
    """Test the logs endpoint."""
    mock_cluster_job_body = mock.MagicMock()
    mock_cluster_job_body.cluster_name = 'test-cluster'
    background_tasks = fastapi.BackgroundTasks()

    # Create an event to track when logs streaming starts
    streaming_started = threading.Event()

    # Mock the stream_response function
    def mock_stream_response(*args, **kwargs):
        streaming_started.set()
        return fastapi.responses.StreamingResponse(
            content=iter([]),  # Empty iterator for testing
            media_type='text/plain')

    def slow_execute(*args, **kwargs):
        # Simulate slow execution
        task = asyncio.create_task(asyncio.sleep(0.1))
        return executor.CoroutineTask(task)

    with mock.patch('sky.server.requests.executor.prepare_request_async') as mock_prepare_async, \
         mock.patch('sky.server.requests.executor.execute_request_in_coroutine',
                   side_effect=slow_execute) as mock_execute, \
         mock.patch('sky.server.stream_utils.stream_response',
                   side_effect=mock_stream_response) as mock_stream:

        # Mock prepare_request to return a request task
        mock_request_task = mock.MagicMock()
        mock_request_task.log_path = '/tmp/test.log'
        mock_prepare_async.return_value = mock_request_task

        # Start logs endpoint in background
        logs_task = asyncio.create_task(
            server.logs(mock.MagicMock(), mock_cluster_job_body,
                        background_tasks))

        # Execute should be run in background and does not block streaming start
        streaming_started.wait(timeout=0.1)

        # Verify the response was created
        response = await logs_task
        assert isinstance(response, fastapi.responses.StreamingResponse)
        assert response.media_type == 'text/plain'
        await background_tasks()

        # Verify the executor calls
        mock_prepare_async.assert_called_once()
        mock_execute.assert_called_once_with(mock_request_task)
        mock_stream.assert_called_once_with(mock.ANY,
                                            mock_request_task.log_path,
                                            mock.ANY,
                                            polling_interval=1,
                                            kill_request_on_disconnect=False)


@mock.patch('sky.utils.context_utils.hijack_sys_attrs')
@mock.patch('asyncio.run')
def test_server_run_uses_uvloop(mock_asyncio_run, mock_hijack_sys_attrs):
    """Test that Server.run uses uvloop event loop policy."""
    from sky.server.uvicorn import Server

    threads_before = len(threading.enumerate())

    config = uvicorn.Config(app='sky.server.server:app',
                            host='127.0.0.1',
                            port=8000)
    server_instance = Server(config)
    original_setup = config.setup_event_loop

    uvloop_policy_set = False
    uvloop_available = True

    def setup_and_check():
        # Save previous event loop policy
        previous_policy = asyncio.get_event_loop_policy()
        try:
            # Call original setup to configure event loop
            original_setup()
            # Check if uvloop policy is now set
            nonlocal uvloop_policy_set, uvloop_available
            try:
                import uvloop
                policy = asyncio.get_event_loop_policy()
                uvloop_policy_set = isinstance(policy, uvloop.EventLoopPolicy)
            except ImportError:
                # uvloop not available
                uvloop_available = False
        finally:
            # Restore previous event loop policy
            # This is needed because other tests/fixtures running on the same
            # pytest worker may not work with the uvicorn event loop policy,
            # such as _seed_test_jobs in test_managed_jobs_service.py
            asyncio.set_event_loop_policy(previous_policy)

    with mock.patch.object(config,
                           'setup_event_loop',
                           side_effect=setup_and_check):
        # Call server.run
        server_instance.run()

    mock_asyncio_run.assert_called_once()
    threads_after = len(threading.enumerate())
    assert threads_after == threads_before

    # Check uvloop policy was set (if uvloop is available)
    if uvloop_available:
        assert uvloop_policy_set, (
            "Expected uvloop event loop policy to be set when uvloop "
            "is available")
    else:
        pytest.skip("uvloop not available, skipping uvloop policy check")


@pytest.mark.asyncio
async def test_enabled_clouds_respect_auth_user():
    auth_user = models.User(id='auth-user-id', name='Auth User')
    request = mock.MagicMock()
    request.state = mock.MagicMock()
    request.state.request_id = 'request-id'
    request.state.auth_user = auth_user

    default_env_vars = {
        constants.USER_ID_ENV_VAR: 'default-id',
        constants.USER_ENV_VAR: 'default-name',
    }

    with mock.patch('sky.server.requests.payloads.request_body_env_vars',
                    side_effect=lambda: default_env_vars.copy()), \
         mock.patch('sky.server.server.executor.schedule_request_async',
                    new_callable=mock.AsyncMock) as mock_schedule:
        await server.enabled_clouds(request, workspace='ws', expand=True)

    mock_schedule.assert_awaited_once()
    _, kwargs = mock_schedule.call_args
    request_body = kwargs['request_body']
    assert request_body.env_vars[constants.USER_ID_ENV_VAR] == auth_user.id
    assert request_body.env_vars[constants.USER_ENV_VAR] == auth_user.name
